package rsautil;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * 加解密工具类，使用RSA加密算法，使用公钥加密、私钥解密
 * 1. 可以通过setPublicKeyStr()、setPrivateKeyStr()两个方法设置公钥私钥字符串
 * 2. 可以通过设置publicKeyStorePath，privateKeyStorePath从class path下的加载指定文件得到公钥私钥串
 * 3. 如果都没有设置，会尝试从class path下的publicKey.keystore，privateKey.keystore两个文件中获取公钥和私钥字符串
 * 使用时应该以单例的方式维护对象，避免为同一对秘钥反复创建对象
 * Created by yaoyunqing on 2020/05/20
 */
public class EncryptionUtil {
    // 公钥字符串
    private String publicKeyStr ;
    // 私钥字符串
    private String privateKeyStr;
    // 默认从ClassPath下取privateKey.keystore文件
    private String publicKeyStorePath = "publicKey.keystore";
    // 默认从ClassPath下取publicKey.keystore文件
    private String privateKeyStorePath = "privateKey.keystore";

    private volatile RSAPublicKey publicKey;
    private volatile RSAPrivateKey privateKey;

    public void setPublicKeyStr(String publicKeyStr) {
        this.publicKeyStr = publicKeyStr;
    }

    public void setPrivateKeyStr(String privateKeyStr) {
        this.privateKeyStr = privateKeyStr;
    }

    public void setPublicKeyStorePath(String publicKeyStorePath) {
        this.publicKeyStorePath = publicKeyStorePath;
    }

    public void setPrivateKeyStorePath(String privateKeyStorePath) {
        this.privateKeyStorePath = privateKeyStorePath;
    }

    /**
     * 加密指定内容
     * @param content
     * @return
     * @throws EncryptionException
     */
    public  String encode(String content) throws EncryptionException {
        loadPublicKey();
        if (publicKey == null) {
            throw new EncryptionException("加密失败，没有可用的公钥");
        }
        return encrypt(publicKey, content);
    }

    /**
     * 解密指定内容
     * @param content
     * @return
     * @throws EncryptionException
     */
    public String decode(String content) throws EncryptionException {
        loadPrivateKey();
        if (privateKey == null) {
            throw new EncryptionException("解密失败，没有可用私钥");
        }
        byte[] decrypt = decrypt(privateKey, Base64Util.decode(content));
        return new String(decrypt, StandardCharsets.UTF_8);
    }


    private void loadPublicKey() throws EncryptionException {
        if (publicKey != null) {
            return;
        }
        synchronized (EncryptionUtil.class) {
            if (publicKey != null) {
                return;
            }
            // 检查 是否定义了字符串
            if (publicKeyStr != null && !publicKeyStr.isEmpty()) {
                loadPublicKeyByStr(publicKeyStr);
            } else if (publicKeyStorePath != null && !publicKeyStorePath.isEmpty()) {
                this.publicKeyStr = getKeyStrFromClassPathFile(publicKeyStorePath);
                if (this.publicKeyStr != null && !this.publicKeyStr.isEmpty()) {
                    loadPublicKeyByStr(this.publicKeyStr);
                }
            }
        }
    }

    /**
     * 从 class path 下的指定文件加载秘钥串
     * @param fileLocation
     * @return
     */
    private String getKeyStrFromClassPathFile(String fileLocation) {
        InputStream in = EncryptionUtil.class.getClassLoader().getResourceAsStream(fileLocation);
        if (in == null) { // class path 下无对应的文件
            return null;
        }
        try (BufferedReader br = new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8))) {
            String line = null;
            while ((line = br.readLine()) != null) {
                if (line.isEmpty()) {
                    continue;
                }
                return line;
            }
            return null;
        } catch (IOException e) {
            e.printStackTrace();
            return null;
        }
    }

    private void loadPrivateKey() throws EncryptionException {
        if (privateKey != null) {
            return;
        }
        synchronized (EncryptionUtil.class) {
            if (privateKey != null) {
                return;
            }
            // 检查 是否定义了字符串
            if (privateKeyStr != null) {
                loadPrivateKeyByStr(privateKeyStr);
            } else if (privateKeyStorePath != null) {
                privateKeyStr = getKeyStrFromClassPathFile(privateKeyStorePath);
                if (privateKeyStr != null && !privateKeyStr.isEmpty()) {
                    loadPrivateKeyByStr(privateKeyStr);
                }
            }
        }
    }

    /**
     * 使用公钥加密
     */
    private String encrypt(RSAPublicKey rsaPublicKey, String content) throws EncryptionException {
        final byte[] encrypt = encrypt(rsaPublicKey, content.getBytes(StandardCharsets.UTF_8));
        // Base64编码后的字符串
        String cipher = Base64Util.encode(encrypt);
        return cipher;
    }


    /**
     * 从字符串中加载公钥
     *
     * @param publicKeyStr 公钥数据字符串
     * @throws EncryptionException 加载公钥时产生的异常
     */
    private void loadPublicKeyByStr(String publicKeyStr) throws EncryptionException {
        if (publicKeyStr == null || publicKeyStr.trim().isEmpty()) {
            throw new EncryptionException("通过字符串加载公钥失败，公钥字符串不能为空");
        }
        try {
            byte[] buffer = Base64Util.decode(publicKeyStr);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(buffer);
            publicKey = (RSAPublicKey) keyFactory.generatePublic(keySpec);
        } catch (NoSuchAlgorithmException e) {
            throw new EncryptionException("通过字符串加载公钥失败，无此算法", e);
        } catch (InvalidKeySpecException e) {
            throw new EncryptionException("通过字符串加载公钥失败，无效的公钥", e);
        } catch (Exception e) {
            throw new EncryptionException("通过字符串加载公钥失败", e);
        }
    }


    /**
     * 从字符串中加载私钥
     *
     * @param privateKeyStr 私钥数据字符串
     * @throws EncryptionException 加载私钥时产生的异常
     */
    private void loadPrivateKeyByStr(String privateKeyStr) throws EncryptionException {
        if (privateKeyStr == null || privateKeyStr.trim().isEmpty()) {
            throw new EncryptionException("通过字符串加载私钥失败，私钥字符串不能为空");
        }
        try {
            byte[] buffer = Base64Util.decode(privateKeyStr);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(buffer);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            privateKey = (RSAPrivateKey) keyFactory.generatePrivate(keySpec);
        } catch (NoSuchAlgorithmException e) {
            throw new EncryptionException("通过字符串加载私钥失败，无此算法", e);
        } catch (InvalidKeySpecException e) {
            throw new EncryptionException("通过字符串加载私钥失败，无效的私钥", e);
        } catch (Exception e) {
            throw new EncryptionException("通过字符串加载私钥失败", e);
        }
    }


    /**
     * 公钥加密过程
     *
     * @param publicKey     公钥
     * @param plainTextData 明文数据
     * @return
     * @throws EncryptionException 加密过程中的异常信息
     */
    private byte[] encrypt(RSAPublicKey publicKey, byte[] plainTextData) throws EncryptionException {
        Cipher cipher = null;
        try {
            // 使用默认RSA
            cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            byte[] output = cipher.doFinal(plainTextData);
            return output;
        } catch (NoSuchAlgorithmException e) {
            throw new EncryptionException("使用公钥加密失败，无此加密算法");
        } catch (InvalidKeyException e) {
            throw new EncryptionException("使用公钥加密失败，加密公钥非法", e);
        } catch (IllegalBlockSizeException e) {
            throw new EncryptionException("使用公钥加密失败，明文长度非法", e);
        } catch (BadPaddingException e) {
            throw new EncryptionException("使用公钥加密失败，明文数据已损坏", e);
        } catch (Exception e) {
            throw new EncryptionException("使用公钥加密失败", e);
        }
    }

    /**
     * 私钥解密过程
     *
     * @param privateKey 私钥
     * @param cipherData 密文数据
     * @return 明文
     * @throws EncryptionException 解密过程中的异常信息
     */
    private  byte[] decrypt(RSAPrivateKey privateKey, byte[] cipherData) throws EncryptionException {
        Cipher cipher = null;
        try {
            // 使用默认RSA
            cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            byte[] output = cipher.doFinal(cipherData);
            return output;
        } catch (NoSuchAlgorithmException e) {
            throw new EncryptionException("使用私钥解密失败，无此解密算法", e);
        } catch (InvalidKeyException e) {
            throw new EncryptionException("使用私钥解密失败，解密私钥非法", e);
        } catch (IllegalBlockSizeException e) {
            throw new EncryptionException("使用私钥解密失败，密文长度非法", e);
        } catch (BadPaddingException e) {
            throw new EncryptionException("使用私钥解密失败，密文数据已损坏", e);
        } catch (Exception e) {
            throw new EncryptionException("密文数据已损坏", e);
        }
    }

}
